import { CCObject, Node, director, log } from "cc";
import Macro from "../Macro";
import { AcademyFixedUpdateStepper } from "./AcademyFixedUpdateStepper";
import { CCCRLCapabilities } from "./Communicator/CCCRLCapabilities";
import { CommunicatorFactory } from "./Communicator/CommunicatorFactory";
import { CommunicatorInitParameters, ICommunicator } from "./Communicator/ICommunicator";
import { ModelRunner } from "./Inference/ModelRunner";
import { RpcCommunicator } from "./Communicator/RpcCommunicator";
import { MLAgentsSettingsManager } from "./MLAgentsSettingsManager";
import { EnvironmentParameters } from "./EnvironmentParameters";
import { StatsRecorder } from "./StatsRecorder";
import { SideChannelManager } from "./SideChannels/SideChannelManager";
import { TrainingAnalyticsSideChannel } from "./SideChannels/TrainingAnalyticsSideChannel";
import { EngineConfigurationChannel } from "./SideChannels/EngineConfigurationChannel";
import { ActionSpec } from "./Actuators/ActionSpec";
import { InferenceDevice } from "./Policies/BarracudaPolicy";
import NNModel from "../Core/Internals/NNModel";

export class Academy {
    // private _environmentParameters: EnvironmentParameters;
    // private _statsRecorder: StatsRecorder;

    /**
     * @en
     * The number of episodes completed by the environment. Incremented
     * each time the environment is reset.
     * 
     * @zh
     * 环境完成的剧集数量。每次重置环境时递增。
     */
    private _episodeCount: number = 0;

    /**
     * @en
     * The number of steps completed in the current episode. Incremented each
     * time a step is taken in the environment. Is reset to 0 during EnvironmentRest.
     * 
     * @zh
     * 当前剧集中完成的步数。每次在环境中执行一步时递增。在环境重置期间重置为0。
     */
    private _stepCount: number = 0;

    /**
     * @en
     * The number of total of steps completed during the whole simulation. Incremented each time
     * a step is taken in the environment.
     * 
     * @zh
     * 在整个模拟过程中完成的总步数。每次在环境中执行一步时递增。
     */
    private _totalStepCount: number = 0;

    /**
     * @en
     * Flag used to keep track fo the first time the Academy is reset.
     * @zh
     * 用于跟踪学院第一次重置的标志。
     */
    private _hadFirstReset: boolean = false;

    /**
     * @en
     * Pointer to the Communicator currently in use by the Academy.
     * 
     * @zh
     * 当前由学院使用的通信器的指针。
     */
    private _communicator: ICommunicator = null!;
    public get communicator(): ICommunicator {
        return this._communicator;
    }
    public set communicator(value: ICommunicator) {
        this._communicator = value;
    }


    private _modelRunners: ModelRunner[] = [];

    /**
     * @en
     * The seed used for inference.
     * 
     * @zh
     * 用于推理的种子。
     */
    private _inferenceSeed: number = 0;
    set inferenceSeed(value: number) {
        this._inferenceSeed = value;
    }

    /**
     * @en
     * Number of training areas to instantiate.
     * 
     * @zh
     * 要实例化的训练区域的数量。
     */
    private _numAreas: number = 0;
    get numAreas(): number {
        return this._numAreas;
    }

    /**
     * @en
     * Returns the RLCapabilities of the python client that the unity process is connected to.
     * 
     * @zh
     * 返回unity进程连接到的python客户端的RLCapabilities。
     */
    private _unityRLCapabilities: CCCRLCapabilities = null!;
    get unityRLCapabilities(): CCCRLCapabilities {
        return this._unityRLCapabilities;
    }
    set unityRLCapabilities(value: CCCRLCapabilities) {
        this._unityRLCapabilities = value;
    }




    private _initialized: boolean = false;
    // private _stepperObject: GameObject;

    static _kApiVersion: string = "1.5.0";
    static _kPackageVersion: string = "2.3.0-exp.3";
    static _KTrainingPort: number = 5004;
    static _portCommandLineFlag = "--mlagents-port";

    /** 
     * @en
     * Reports whether or not the communicator is on.
     * @zh
     * 报告通信器是否打开。
     */
    get isCommunicatorOn(): boolean {
        return this._communicator != null;
    }


    static get isInitialized(): boolean {
        return this._instance != null;
    }

    private _trainerCapabilities: CCCRLCapabilities = null!;
    public get trainerCapabilities(): CCCRLCapabilities {
        return this._trainerCapabilities;
    }
    public set trainerCapabilities(value: CCCRLCapabilities) {
        this._trainerCapabilities = value;
    }


    /**
     * @en
     * The Academy uses a series of events to communicate with agents
        to facilitate synchronization. More specifically, it ensures
        that all the agents perform their steps in a consistent order (i.e. no
        agent can act based on a decision before another agent has had a chance
        to request a decision).
        
        @zh
        学院使用一系列事件与代理进行通信，以促进同步。
        更具体地说，它确保所有代理以一致的顺序执行其步骤（即在另一个代理请求决策之前，没有代理可以根据决策采取行动）。
     */


    /**
     * 
     * @en
     * Signals to all the Agents at each environment step so they can use
        their Policy to decide on their next action.

        @zh
        在每个环境步骤中向所有代理发出信号，以便他们可以使用
     */
    decideAction: () => void = null!;


    /**
     * @en
     * Signals to all the listeners that the academy is being destroyed
     * 
     * @zh
     * 向所有侦听器发出信号，学院正在被销毁
     */
    destroyAction: () => void = null!;


    /**
     * @en
     * Signals to the Agent that a new step is about to start.
        This will mark the Agent as Done if it has reached its maxSteps.

        @zh
        向代理发出信号，即将开始新的步骤。
     */
    agentIncrementStep: () => void = null!;



    /**
     * @en
     * Signals to all of the <see cref="Agent"/>s that their step is about to begin.
       This is a good time for an <see cref="Agent"/> to decide if it would like to
       call <see cref="Agent.RequestDecision"/> or <see cref="Agent.RequestAction"/>
       for this step.  Any other pre-step setup could be done during this event as well.


       @zh
        向所有<see cref="Agent"/>发出信号，即将开始他们的步骤。
        这是一个很好的时机，<see cref="Agent"/>可以决定是否要
        调用<see cref="Agent.RequestDecision"/>或<see cref="Agent.RequestAction"/>
        为这一步。 任何其他预步骤设置也可以在此事件期间完成。
     */
    agentPreStep: (i: number) => void = null!;


    /**
     * @en
     * Signals to all the agents each enviroment step so they can shend their state to their Policy if they have requested a decision.
     * 
     * @zh
     * 在每个环境步骤中向所有代理发出信号，以便他们可以将状态发送到他们的策略，如果他们请求了一个决定。
     */
    agentSendState: () => void = null!;


    /**
     * @en
     * Signals to all the agents at each environment step so they can act if
       they have requested a decision.

       @zh
        在每个环境步骤中向所有代理发出信号，以便他们可以行动如果他们请求了一个决定。
     */
    agentAct: () => void = null!;

    /**
     * @en
     * Signals to all the agents each time the Academy force resets.
     * 
     * @zh
     * 每次学院强制重置时向所有代理发出信号。
     */
    agentForceReset: () => void = null!;

    /**
     * @en
     * signals that the Academy has been reset by the training process.
     * 
     * @zh
     * 信号表明学院已被训练过程重置。
     */
    onEnvironmentReset: () => void = null!;

    _fixedUpdateStepper: AcademyFixedUpdateStepper = null!;

    _stepperObject: Node = null!;

    private constructor() {
        if (Macro.CCC_DEV) {
            if (!CommunicatorFactory.communicatorRegistered) {
                log("CommunicatorFactory not registered");
                CommunicatorFactory.register(RpcCommunicator.create);
            }
        }

        this.lazyInitialize();
    }

    private static _instance: Academy = null!;
    public static get instance(): Academy {
        if (this._instance == null) {
            this._instance = new Academy();
        }
        return this._instance;
    }


    lazyInitialize(): void {
        if (!this._initialized) {
            this.initializeEnvironment();
            this._initialized = true;
        }
    }

    /**
     * @en 
     *  Enable stepping of the Academy during the FixedUpdate phase. This is done by creating
        a temporary Node with a Component that calls Academy.EnvironmentStep().

        @zh
        在FixedUpdate阶段启用学院的步进。这是通过创建
        一个临时Node，其中包含一个Component，该Component调用Academy.EnvironmentStep()。
     * @returns 
     */
    enableAutomaticStepping() {
        if (this._fixedUpdateStepper != null) {
            return;
        }

        this._stepperObject = new Node("AcademyFixedUpdateStepper");
        this._stepperObject.hideFlags = CCObject.Flags.HideInHierarchy;
        this._fixedUpdateStepper = this._stepperObject.addComponent(AcademyFixedUpdateStepper);
        this._stepperObject.parent = director.getScene();

    }

    disableAutomaticStepping(): void {
        if (this._fixedUpdateStepper == null) {
            return;
        }

        this._fixedUpdateStepper.node.destroy();

        this._fixedUpdateStepper = null!;

        this._stepperObject = null!;
    }

    get automaticSteppingEnabled(): boolean {
        return this._fixedUpdateStepper != null;
    }

    set automaticSteppingEnabled(value: boolean) {
        if (value) {
            this.enableAutomaticStepping();
        } else {
            this.disableAutomaticStepping();
        }
    }


    static readPortFromArgs(): number {
        if (Macro.CCC_DEV) {
            return MLAgentsSettingsManager.settings.connectTrainer ? MLAgentsSettingsManager.settings.editorTrainingPort : -1;
        }

        return -1;
    }


    _environmentParameters: EnvironmentParameters = null!;
    /**
     * @en
     * Returns the {@link EnvironmentParameters} instance. If training
     * features such as Curriculum Learning or Environment Parameter Randomization are used,
     * then the values of the parameters generated from the training process can be retrieved here.
     * 
     * @zh
     * 返回{@link EnvironmentParameters}实例。如果使用了训练功能，例如课程学习或环境参数随机化，
     * 那么可以在此处检索从训练过程中生成的参数的值。
     * 
     */
    public get environmentParameters(): EnvironmentParameters {
        return this._environmentParameters;
    }

    _statsRecorder: StatsRecorder = null!;

    /**
     * @en
     *  Returns the {@link StatsRecorder} instance. This instance can be used
     * to record any statistics from the CCC environment.
     * 
     * @zh
     * 返回{@link StatsRecorder}实例。此实例可用于记录来自CCC环境的任何统计信息。
     */
    public get statsRecorder(): StatsRecorder {
        return this._statsRecorder;
    }


    public async initializeEnvironment(): Promise<void> {

        this.enableAutomaticStepping();


        SideChannelManager.registerSideChannel(new EngineConfigurationChannel());
        SideChannelManager.registerSideChannel(new TrainingAnalyticsSideChannel());
        this._environmentParameters = new EnvironmentParameters();
        this._statsRecorder = new StatsRecorder();
        const port = Academy.readPortFromArgs();
        if (port > 0) {
            const tempCommunicator = CommunicatorFactory.create();
            if (tempCommunicator) {
                this._communicator = tempCommunicator;
            }
        }

        if (!this._communicator && CommunicatorFactory.enabled && port > 0) {
            log("Communicator failed to start!");
        }

        if (this._communicator) {

            /**
             * 
             * @en
             * We try to exchange the first message with Python. If this fails, it means
             *  no Python Process is ready to train the environment. In this case, the
             * environment must use Inference.
             */
            /**
             * @zh
             * 我们尝试与Python交换第一条消息。如果失败，这意味着
             * 没有Python进程准备训练环境。在这种情况下，环境必须使用推理。
             */

            const communicatorInitParams: CommunicatorInitParameters = {
                port: port,
                communicationVersion: Academy._kApiVersion,
                packageVersion: Academy._kPackageVersion,
                name: "AcademySingleton",
                cccCapabilities: new CCCRLCapabilities()
            };

            try {
                const [initSuccessful, unityRlInitParameters] = await this._communicator.initialize(
                    communicatorInitParams
                );

                if (initSuccessful) {
                    this._inferenceSeed = unityRlInitParameters.seed;
                    this._numAreas = unityRlInitParameters.numAreas;
                    this.trainerCapabilities = unityRlInitParameters.trainerCapabilities;
                    this.trainerCapabilities.warnOnPythonMissingBaseRLCapabilities();
                    log("1 初始化环境成功 ");
                } else {
                    log(`Couldn't connect to trainer on port ${port} using API version ${Academy._kApiVersion}. Will perform inference instead.`);
                    this._communicator = null!;
                }


            } catch (error) {
                console.error(error);
            }
        }

        if (this._communicator != null) {
            this._communicator.quitCommandReceived = this.onQuitCommandReceived;
            this._communicator.resetCommandReceived = this.onResetCommand;
        }

        this.resetActions();
    }


    private resetActions(): void {
        // this.decideAction = () => { };
        // this.destroyAction = () => { };
        // this.agentPreStep = (i: number) => { };
        // this.agentSendState = () => { };
        // this.agentAct = () => { };
        // this.agentForceReset = () => { };
        // this.onEnvironmentReset = () => { };
    }


    private onQuitCommandReceived(): void {

    }

    onResetCommand(): void {
        this.forcedFullReset();
    }

    public get episodeCount(): number {
        return this._episodeCount;
    }



    public get stepCount(): number {
        return this._stepCount;
    }

    public get totalStepCount(): number {
        return this._totalStepCount;
    }

    // public get communicator(): Communicator | null {
    //     return this._communicator;
    // }


    private forcedFullReset(): void {
        this.environmentReset();
        if (this.agentForceReset) {
            this.agentForceReset();
        }
        this._hadFirstReset = true;
    }

    /**
     * 
     * @en
     *  Performs a single environment update of the Academy and Agent
     * 
     * @zh
     * 执行学院和代理的单个环境更新
     */
    public environmentStep(): void {
        if(!this._initialized){
            return;
        }

        if (!this._hadFirstReset) {
            this.forcedFullReset();
        }

        if (this.agentPreStep) {
            this.agentPreStep(this._stepCount);
        }
        this._stepCount += 1;
        this._totalStepCount += 1;
        if (this.agentIncrementStep) {
            this.agentIncrementStep();
        }
        if (this.agentSendState) {
            this.agentSendState();
        }
        if (this.decideAction) {
            this.decideAction();
        }

        if (!this.isCommunicatorOn) {
            SideChannelManager.getSideChannelMessage();
        }

        if (this.agentAct) {
            this.agentAct();
        }
    }

    private environmentReset(): void {
        this._stepCount = 0;
        this._episodeCount++;
        if (this.onEnvironmentReset) {
            this.onEnvironmentReset();
        }
    }

    public getOrCreateModelRunner(
        model: NNModel,
        actionSpec: ActionSpec,
        inferenceDevice: InferenceDevice,
        deterministicInference: boolean = false
    ): ModelRunner {
        let modelRunner = this._modelRunners.find(x => x.hasModel(model, inferenceDevice));
        if (!modelRunner) {
            modelRunner = new ModelRunner(model, actionSpec, inferenceDevice, this._inferenceSeed, deterministicInference);
            this._modelRunners.push(modelRunner);
            this._inferenceSeed++;
        }
        return modelRunner;
    }

    // public dispose(): void {
    //     this.disableAutomaticStepping();
    //     if (this.destroyAction) {
    //         this.destroyAction();
    //     }
    //     if (this._communicator) {
    //         this._communicator.dispose();
    //         this._communicator = null;
    //     }
    //     this._environmentParameters.dispose();
    //     this._statsRecorder.dispose();
    //     SideChannelManager.unregisterAllSideChannels();
    //     if (this._modelRunners) {
    //         for (const modelRunner of this._modelRunners) {
    //             modelRunner.dispose();
    //         }
    //         this._modelRunners = null;
    //     }
    //     this.resetActions();
    //     TimerStack.instance.saveJsonTimers();
    //     this._initialized = false;
    // }




    // private get isCommunicatorOn(): boolean {
    //     // TODO: Implement the logic to check if the communicator is on
    //     return false;
    // }

    isStepperOwner(stepper: AcademyFixedUpdateStepper) {
        if (stepper.node === Academy.instance._stepperObject) {
            return true;
        } else {
            return false;
        }
    }
}